from ultralytics import YOLO
import cv2

def predict_image_model(model_path, image_path, output_path=None):
    # 加载模型
    model = YOLO(model_path)
    
    # 预测：使用与训练一致的参数
    results = model(
        image_path,
        imgsz=640,               # 与训练时输入尺寸一致
        conf=0.25,               # 训练时默认置信度阈值
        iou=0.7,                 # 与训练时NMS IOU一致
        rect=False,              # 关闭矩形推理，与训练逻辑匹配
        augment=False,           # 关闭推理增强
        verbose=False
    )
    
    # 可视化结果（使用兼容所有版本的参数）
    for r in results:
        im_array = r.plot(
            conf=True,           # 显示置信度
            boxes=True,          # 显示检测框
            kpt_line=True        # 显示关键点连接线（注意参数是kpt_line而非kptline）
        )
        im = cv2.cvtColor(im_array, cv2.COLOR_RGB2BGR)  # 转换为BGR格式
        
        # 显示图像
        cv2.imshow('Result', im)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
        
        # 保存结果
        if output_path:
            cv2.imwrite(output_path, im)
            print(f"结果已保存到: {output_path}")

if __name__ == "__main__":
    # 模型路径（根据实际训练结果目录调整）
    model_path = "model/last.pt"
    # 测试图片路径
    image_path = "test_card.jpg"
    # 输出结果保存路径
    output_path = "result.jpg"
    
    # 调用函数时使用正确的函数名（之前的代码存在函数名拼写错误）
    predict_image_model(model_path, image_path, output_path)
